# -*- coding: utf-8 -*-
# Thanks to the awesome tools:
# - Windows Exploit Suggester - Next Generation: https://github.com/bitsadmin/wesng
# - linux-exploit-suggester: https://github.com/mzet-/linux-exploit-suggester

import os
import imp
import threading

from urllib2 import urlopen
from zipfile import BadZipfile

from pupylib import ROOT
from pupylib.PupyOutput import List, Table, Color, MultiPart, NewLine
from pupylib.PupyModule import config, PupyModule, PupyArgumentParser

__class_name__ = 'Exploit_Suggester'

LINUX_EXPLOIT_SUGGESTER_PATH = os.path.join(
    ROOT, 'external', 'linux-exploit-suggester', 'linux-exploit-suggester.sh'
)

WES_PATH = os.path.join(
    ROOT, 'external', 'wesng', 'wes.py'
)

WES_DEFINITIONS = 'https://raw.githubusercontent.com/bitsadmin/wesng/master/definitions.zip'
WES_LOCAL_FILE = 'wes-defintions.zip'


@config(compat=['linux', 'windows'], category='exploit')
class Exploit_Suggester(PupyModule):
    ''' Exploit suggester '''

    terminate_pipe = None
    terminated = False

    dependencies = ['pupyutils.safepopen']

    @classmethod
    def init_argparse(cls):
        cls.arg_parser = PupyArgumentParser(prog='Exploit_Suggester', description=cls.__doc__)
        cls.arg_parser.add_argument(
            '-no-recent-kb', default=False, action='store_true',
            help='Do not filter findings by most recent KB date')
        cls.arg_parser.add_argument(
            '--hide', nargs='+', default='',
            help='(WES only) Hide vulnerabilities of for example Adobe Flash Player and Microsoft Edge')
        cls.arg_parser.add_argument(
            '--update', action='store_true', default=False,
            help='(WES only) Update Windows database (Internet access required on pupy server host)')

    def run(self, args):
        close_event = threading.Event()
        result = []

        cmdargs = None
        kwargs = None

        definitions = None

        safe_exec = self.client.remote('pupyutils.safepopen', 'safe_exec', False)

        if self.client.is_linux():
            payload = open(LINUX_EXPLOIT_SUGGESTER_PATH).read()
            cmdargs = ['/bin/bash']
            kwargs = (('stdin_data', payload),)
        else:
            definitions = os.path.join(
                self.config.get_folder('plugins'), WES_LOCAL_FILE
            )

            if not os.path.isfile(definitions) or args.update:
                self.info('Updating WES defintions from {}'.format(WES_DEFINITIONS))
                try:
                    response = urlopen(WES_DEFINITIONS)
                    with open(definitions, 'w+b') as out:
                        while True:
                            block = response.read(32768)
                            if not block:
                                break

                            out.write(block)

                except Exception as e:
                    self.error('Update failed: {}'.format(e))

                    if os.path.isfile(definitions):
                        try:
                            os.unlink(definitions)
                        except (OSError, IOError):
                            pass

                    return

                self.info('Update completed ({})'.format(definitions))

            expandvars = self.client.remote('os.path', 'expandvars')
            systeminfo = expandvars(r'%WINDIR%\System32\systeminfo.exe')
            cmdargs = [systeminfo]
            kwargs = tuple()

        self.info('Execute payload ({})'.format(' '.join(cmdargs)))
        self.terminate_pipe, get_returncode = safe_exec(
            result.append, close_event.set, cmdargs, kwargs
        )

        close_event.wait()
        retcode = get_returncode()

        if retcode != 0:
            self.warning('Ret: {}'.format(retcode))
        else:
            self.success('Done')

        result = ''.join(result)

        if not result:
            self.error('No data')
            return

        if self.client.is_linux():
            self.log(result)
            return

        wes = imp.load_source('wes', WES_PATH)

        try:
            cves, date = wes.load_definitions(definitions)
        except BadZipfile:
            self.error(
                'Defintions were downloaded incorrectly ({})'.format(
                    definitions))
            return

        productfilter, win, mybuild, version, arch, hotfixes = \
          wes.determine_product(result)

        self.log(
            List([
                'Definitions: ' + str(date),
                'Name: ' + productfilter,
                'Generation: ' + (win or 'N/A'),
                'Build: ' + (str(mybuild) if mybuild else 'N/A'),
                'Version: ' + (str(version) or 'N/A'),
                'Architecture: ' + arch,
                'Hotfixes: ' + ', '.join([
                    'KB%s' % kb for kb in hotfixes
                ])
            ], caption='Operating System')
        )

        try:
            filtered, found = wes.determine_missing_patches(
                productfilter, cves, hotfixes)
        except wes.WesException as e:
            self.error(e.msg)
            return

        if not args.no_recent_kb:
            recentkb = wes.get_most_recent_kb(found)
            if recentkb:
                recentdate = int(recentkb['DatePosted'])
                found = list(filter(lambda kb: int(kb['DatePosted']) >= recentdate, found))

        if 'Windows Server' in productfilter:
            self.info('Filtering duplicate vulnerabilities')
            found = wes.filter_duplicates(found)

        filtered = wes.apply_display_filters(found, args.hide, True, [], [])
        if not filtered:
            self.info('No vulnerabilities found')
            return

        results = {}
        proposed = set()

        for res in filtered:
            exploits = res['Exploits'].split(',')
            for exploit in exploits:
                exploit = exploit.strip()
                if exploit in proposed:
                    continue

                proposed.add(exploit)

                impact = ''.join(part[0] for part in res['Impact'].split())
                color = 'white'
                if impact == 'ID':
                    color = 'grey'
                elif res['Severity'] == 'Critical' or impact in ('RCE', 'EoP'):
                    color = 'lightred'
                elif res['Severity'] == 'Important':
                    color = 'lightyellow'

                title = (res['AffectedComponent'] + ' / ' + res['AffectedProduct']) \
                  if res['AffectedComponent'] else res['AffectedProduct']

                if title not in results:
                    results[title] = []

                results[title].append({
                    'CVE': Color(res['CVE'], color),
                    'Date': res['DatePosted'],
                    'Impact': impact,
                    'Exploit': exploit
                })

        tables = [NewLine()]
        for component, cves in results.iteritems():
            tables.append(
                Table(cves, ['CVE', 'Date', 'Impact', 'Exploit'], component))

        self.log(MultiPart(tables))

    def interrupt(self):
        if not self.terminated and self.terminate_pipe:
            self.terminated = True
            self.error('Stopping command')
            self.terminate_pipe()
            self.error('Stopped')
